import dataclasses
import io
import mmap
import os
import struct
import typing
from os import PathLike
from pathlib import Path
from typing import AnyStr, Tuple, BinaryIO, ContextManager, Dict, Type

from cachetools import LRUCache
from typing_extensions import Buffer

from storage.blobstore.base import AbstractBlobStore
from storage.blobstore.compression import CompressMethod, decompress_blob, compress_blob
from storage.blobstore.crypto import EncryptionAlgorithm, decrypt_blob, encrypt_blob
from storage.kvstore.base import AbstractKVStore
from storage.mutex import Mutex, MutexProtocol


@dataclasses.dataclass
class FileBlobHeader:
    """
    FileBlob file header

    +---------+-------------------+-----------------+---------------+------------------+
    | magic   | encrypt algorithm | compress method | dirty length  | available length |
    +---------+-------------------+-----------------+---------------+------------------+
    | char[4] | uint16            | uint16          | uint32        | uint32           |
    +---------+-------------------+-----------------+---------------+------------------+
    """
    MAGIC: typing.ClassVar[bytes] = b'BLOB'
    FMT: typing.ClassVar[str] = '>4sHHII'
    SIZE: typing.ClassVar[int] = struct.calcsize(FMT)

    magic: bytes
    encrypt_algorithm: EncryptionAlgorithm
    compress_method: CompressMethod
    dirty_length: int
    available_length: int

    @classmethod
    def from_bytes(cls, data: Buffer) -> 'FileBlobHeader':
        magic, encryption_algorithm, compress_method, dirty_length, available_length = struct.unpack(cls.FMT, data)
        return cls(magic=magic,
                   encrypt_algorithm=EncryptionAlgorithm(encryption_algorithm),
                   compress_method=CompressMethod(compress_method),
                   dirty_length=dirty_length,
                   available_length=available_length)

    def to_bytes(self) -> bytes:
        return struct.pack(self.FMT,
                           self.magic,
                           self.encrypt_algorithm,
                           self.compress_method,
                           self.dirty_length,
                           self.available_length)


@dataclasses.dataclass
class BlobHeader:
    """
    BlobHeader format

    +-------------+----------+
    | blob length | nonce_iv |
    +-------------+----------+
    | uint32      | char[16] |
    +-------------+----------+
    """
    FMT: typing.ClassVar[str] = '>L16s'
    SIZE: typing.ClassVar[int] = struct.calcsize(FMT)

    blob_length: int
    nonce_iv: bytes

    def to_bytes(self) -> bytes:
        return struct.pack(self.FMT, self.blob_length, self.nonce_iv)

    @classmethod
    def from_bytes(cls, data: Buffer) -> 'BlobHeader':
        [length, nonce_iv] = struct.unpack(cls.FMT, data)
        return cls(blob_length=length, nonce_iv=nonce_iv)


class FileBlobStore(AbstractBlobStore):
    def __init__(self,
                 path: AnyStr | PathLike,
                 kv_store: AbstractKVStore,
                 *,
                 encrypt: EncryptionAlgorithm | int = EncryptionAlgorithm.AES128CTR,
                 secret: bytes = None,
                 compress: CompressMethod | int = CompressMethod.ZSTD,
                 mutex_type=Mutex,
                 batch: int = 100,
                 cache_size=1024,
                 preallocate=10 * 1024 * 1024):
        self._path = Path(path)
        self._fd = os.open(self._path, os.O_RDWR | os.O_CREAT, 0o644)
        if os.fstat(self._fd).st_size == 0:
            os.ftruncate(self._fd, preallocate + FileBlobHeader.SIZE)
            header = FileBlobHeader(magic=FileBlobHeader.MAGIC,
                                    encrypt_algorithm=encrypt,
                                    compress_method=compress,
                                    dirty_length=0,
                                    available_length=preallocate)
            os.write(self._fd, header.to_bytes())
        self._mm: mmap.mmap = mmap.mmap(self._fd, 0, access=mmap.ACCESS_WRITE)
        self._file_blob_header = FileBlobHeader.from_bytes(self._mm[:FileBlobHeader.SIZE])

        self._lock_name = Path(str(self._path.resolve()) + '.lock')

        self._kv_store = kv_store
        self._secret: bytes | None = secret
        if self._secret is None and self._encrypt_algorithm != EncryptionAlgorithm.Plaintext:
            raise ValueError(f'secret must be set when encryption algorithm is not Plaintext')

        self._cache = LRUCache(cache_size)
        self._mutex_type: Type[MutexProtocol] | None = mutex_type

        # 未压缩、未加密、未落盘的待写入数据
        self._put_blob_buffer_size = batch
        self._put_blob_buffer: Dict[bytes, bytes] = dict()

    @property
    def _encrypt_algorithm(self) -> EncryptionAlgorithm:
        return self._file_blob_header.encrypt_algorithm

    @property
    def _compress_method(self) -> CompressMethod:
        return self._file_blob_header.compress_method

    @property
    def _available_length(self) -> int:
        return self._file_blob_header.available_length

    @_available_length.setter
    def _available_length(self, value: int):
        self._file_blob_header.available_length = value

    @property
    def _dirty_length(self) -> int:
        return self._file_blob_header.dirty_length

    @_dirty_length.setter
    def _dirty_length(self, value: int):
        self._file_blob_header.dirty_length = value

    def close(self):
        self._mm[:FileBlobHeader.SIZE] = self._file_blob_header.to_bytes()
        self.flush()
        self._mm.close()
        os.close(self._fd)

    def _locate_blob(self, key: bytes) -> int | None:
        if (data := self._kv_store.get(key)) is None:
            return None
        return int.from_bytes(data, 'big')

    def _read_raw_blob(self, offset: int) -> Tuple[BlobHeader, bytes]:
        header_start = offset
        header_end = offset + BlobHeader.SIZE
        h = BlobHeader.from_bytes(self._mm[header_start:header_end])

        blob_start = offset + BlobHeader.SIZE
        blob_end = blob_start + h.blob_length
        return h, self._mm[blob_start:blob_end]

    def _grow_file(self, length=10 * 1024 * 1024):
        target = os.fstat(self._fd).st_size + length
        self._mm.close()
        os.ftruncate(self._fd, target)
        self._mm = mmap.mmap(self._fd, 0, access=mmap.ACCESS_WRITE)
        self._available_length += length
        self._mm[:FileBlobHeader.SIZE] = self._file_blob_header.to_bytes()

    def _write_raw_blob(self, header: BlobHeader, data: bytes) -> int:
        while self._available_length <= BlobHeader.SIZE + header.blob_length:
            self._grow_file()

        write_start = FileBlobHeader.SIZE + self._dirty_length
        write_end = write_start + BlobHeader.SIZE + header.blob_length
        self._mm[write_start:write_end] = header.to_bytes() + data
        self._dirty_length += BlobHeader.SIZE + header.blob_length
        self._available_length -= BlobHeader.SIZE + header.blob_length
        return write_start

    def has_blob(self, key: bytes) -> bool:
        if key in self._cache:
            return True

        return self._locate_blob(key) is not None

    def get_blob(self, key: bytes) -> bytes | None:
        cached_result = self._cache.get(key)
        if cached_result is not None:
            return cached_result

        if key in self._put_blob_buffer:
            result = self._put_blob_buffer[key]
            self._cache[key] = result
            return result

        if (offset := self._locate_blob(key)) is None:
            return None

        if self._mutex_type is not None:
            with self._mutex_type(str(self._lock_name)):
                header, ciphertext_data = self._read_raw_blob(offset)
        else:
            header, ciphertext_data = self._read_raw_blob(offset)

        plaintext_data = decrypt_blob(self._encrypt_algorithm, self._secret, header.nonce_iv, ciphertext_data)
        decompressed_data = decompress_blob(self._compress_method, plaintext_data)

        self._cache[key] = decompressed_data
        return decompressed_data

    def put_blob(self, key: bytes, data: bytes):
        if len(self._put_blob_buffer) < self._put_blob_buffer_size:
            self._put_blob_buffer[key] = data
            return

        self._flush_to_disk()
        self._put_blob_buffer[key] = data

    def _flush_to_disk(self):
        kv_pairs = dict()
        for key, data in self._put_blob_buffer.items():
            nonce_iv = b''
            encrypt_alg = self._encrypt_algorithm
            if encrypt_alg != EncryptionAlgorithm.Plaintext:
                nonce_iv = os.urandom(16)

            compress_method = self._compress_method

            compressed_data = compress_blob(compress_method, data)
            ciphertext_data = encrypt_blob(encrypt_alg, self._secret, nonce_iv, compressed_data)

            header = BlobHeader(blob_length=len(ciphertext_data), nonce_iv=nonce_iv)

            if self._mutex_type is not None:
                with self._mutex_type(str(self._lock_name)):
                    offset = self._write_raw_blob(header, ciphertext_data)
            else:
                offset = self._write_raw_blob(header, ciphertext_data)

            kv_pairs[key] = offset.to_bytes(8, 'big')

        if len(kv_pairs) > 0:
            self._kv_store.update(kv_pairs)
        self._put_blob_buffer = dict()
        self._mm.flush()

    def flush(self):
        self._flush_to_disk()
